import torch
from transformers import CLIPModel, CLIPProcessor, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig, PreTrainedTokenizer
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
import torch.nn.functional as F
from model_old.model_ChineseCLIPSimilarityCalculator import * # 假设这个模块可用


# 1. 加载 CLIP 模型 (保持不变)
def build_clip_model(clip_model_path):
    processor = CLIPProcessor.from_pretrained(clip_model_path, from_slow=True, use_fast=False)
    model = CLIPModel.from_pretrained(clip_model_path)
    return model, processor

# 2. 加载 LLM（如 Baichuan）并开启 4bit 量化 (保持不变)
def build_llm_model(llm_path):
    tokenizer = AutoTokenizer.from_pretrained(llm_path, use_fast=False, trust_remote_code=True)
    if hasattr(tokenizer, "apply_chat_template"):
        def identity_apply_chat_template(*args, **kwargs):
            if isinstance(args[0], list):
                return args[0][0]
            return args[0]
        tokenizer.apply_chat_template = identity_apply_chat_template

    llm_config = AutoConfig.from_pretrained(llm_path, trust_remote_code=True)
    llm_config.max_position_embeddings = 4096 
    print(f"DEBUG: Setting LLM max_position_embeddings to {llm_config.max_position_embeddings}") # 添加此调试打印以确认！
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    model = AutoModelForCausalLM.from_pretrained(
        llm_path,
        config=llm_config,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True
    )

    if hasattr(model, 'enable_input_require_grads'):
        model.enable_input_require_grads()
    if hasattr(model, 'gradient_checkpointing_enable'):
        model.gradient_checkpointing_enable()
    else:
        print("WARNING: Model does not have gradient_checkpointing_enable method.")
    model = prepare_model_for_kbit_training(model)
    return model, tokenizer

def add_lora_to_models(clip_model, llm_model):
    # --- 冻结所有模型参数 ---
    for param in clip_model.parameters():
        param.requires_grad = False
    for param in llm_model.parameters():
        param.requires_grad = False

    # --- CLIP 配置：解冻最后三层并应用LoRA ---
    clip_global_target_modules = []
    
    if hasattr(clip_model, 'vision_model') and \
       hasattr(clip_model.vision_model, 'encoder') and \
       hasattr(clip_model.vision_model.encoder, 'layers'):
        
        vision_layers = clip_model.vision_model.encoder.layers
        num_clip_layers = len(vision_layers)
        
        if num_clip_layers >= 3:
            last_three_layers = [num_clip_layers - 3, num_clip_layers - 2, num_clip_layers - 1]
            for layer_idx in last_three_layers:
                for param in vision_layers[layer_idx].parameters():
                    param.requires_grad = True
            print(f"[INFO] Unfroze last three layers of CLIP vision encoder (layers {last_three_layers})")
            
            target_proj_names = ["q_proj", "k_proj", "v_proj", "out_proj"]
            target_module_types = (torch.nn.Linear,)
            
            clip_global_target_modules = []
            for name, module in clip_model.named_modules():
                if isinstance(module, target_module_types):
                    if "vision_model.encoder.layers." in name:
                        layer_part = name.split("vision_model.encoder.layers.")[1]
                        layer_idx = int(layer_part.split(".")[0])
                        if layer_idx in last_three_layers:
                            if any(proj in name for proj in target_proj_names):
                                clip_global_target_modules.append(name)
            
            if clip_global_target_modules:
                clip_lora_config = LoraConfig(
                    r=8, 
                    lora_alpha=16,
                    lora_dropout=0.2,
                    bias="none",
                    task_type=TaskType.FEATURE_EXTRACTION,
                    target_modules=clip_global_target_modules,
                    fan_in_fan_out=False 
                )
                clip_model = get_peft_model(clip_model, clip_lora_config)
                print("--- CLIP 可训练参数 ---")
                clip_model.print_trainable_parameters()
            else:
                print("[WARNING] CLIP LoRA未找到合适的target modules。")
        else:
            print(f"[WARNING] CLIP视觉模型只有{num_clip_layers}层（需要至少3层）。")
    else:
        print("[WARNING] CLIP模型结构与预期格式不符。")

    # --- Qwen1.5 LoRA 配置优化 ---
    qwen_lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.2,  
        bias="none",
        task_type="CAUSAL_LM"
    )

    def check_lora_modules(model, target_modules):
        existing_modules = []
        for name, _ in model.named_modules():
            for module in target_modules:
                if name.endswith(module):
                    existing_modules.append(name)
        return existing_modules

    found_modules = check_lora_modules(llm_model, qwen_lora_config.target_modules)
    
    print("--- LLM 可训练参数 (LoRA之前) ---")
    total_params = sum(p.numel() for p in llm_model.parameters())
    trainable_params = sum(p.numel() for p in llm_model.parameters() if p.requires_grad)
    print(f"可训练: {trainable_params} ({trainable_params/total_params*100:.4f}%) | 总数: {total_params}")
    
    if found_modules:
        llm_model = get_peft_model(llm_model, qwen_lora_config)
        print("--- LLM 可训练参数 (LoRA之后) ---")
        llm_model.print_trainable_parameters()
    else:
        print("[WARNING] Qwen1.5中未找到合适的target modules。")
    
    return clip_model, llm_model   

class MultiModalContrastiveLoss(torch.nn.Module):
    """图像-文本对比学习损失（适配CLIP与LLM特征）"""
    def __init__(self, temperature):
        super().__init__()
        self.temperature = temperature
        self.cross_entropy = torch.nn.CrossEntropyLoss()

    def forward(self, image_features, text_features):
        image_features = F.normalize(image_features, dim=-1)
        text_features = F.normalize(text_features, dim=-1)

        logits = torch.matmul(image_features, text_features.t()) / self.temperature
        batch_size = logits.shape[0]
        labels = torch.arange(batch_size, device=logits.device)

        loss_i2t = self.cross_entropy(logits, labels)
        loss_t2i = self.cross_entropy(logits.t(), labels)
        contrastive_loss = (loss_i2t + loss_t2i) / 2

        with torch.no_grad():
            acc_i2t = (logits.argmax(dim=1) == labels).float().mean()
            acc_t2i = (logits.t().argmax(dim=1) == labels).float().mean()

        return contrastive_loss, {"acc_i2t": acc_i2t, "acc_t2i": acc_t2i}


class MultiModalTaggingModel(torch.nn.Module):
    def __init__(self, clip_model, llm_model, processor, tokenizer, similarity_calculator=None, 
                 contrastive_loss_weight=0.5, contrastive_genloss_weight=0.4):
        super().__init__()
        self.clip = clip_model
        self.llm = llm_model
        self.processor = processor  # CLIP的processor
        self.tokenizer = tokenizer  # LLM的tokenizer
        
        self.similarity_calculator = similarity_calculator
        self.is_preprocessed = False
        self.similarity_threshold = 0.35
        
        vision_hidden_size = self.clip.vision_model.config.hidden_size
        d_llm = llm_model.config.hidden_size
        self.clip2llm_proj = torch.nn.Sequential(
            torch.nn.Linear(vision_hidden_size, d_llm),
            torch.nn.Dropout(0.2)
        )
        if isinstance(self.clip2llm_proj[0], torch.nn.Linear):
            torch.nn.init.xavier_normal_(self.clip2llm_proj[0].weight)
            if self.clip2llm_proj[0].bias is not None:
                torch.nn.init.zeros_(self.clip2llm_proj[0].bias)
        for param in self.clip2llm_proj.parameters():
            param.requires_grad = True
        
        # ---------------- 新增：对比学习相关组件 ----------------
        contrastive_dim = 512
        self.image_contrast_proj = torch.nn.Sequential(
            torch.nn.Linear(vision_hidden_size, contrastive_dim),
            torch.nn.BatchNorm1d(contrastive_dim),
            torch.nn.GELU(),
            torch.nn.Linear(contrastive_dim, contrastive_dim),
            torch.nn.BatchNorm1d(contrastive_dim)
        )
        for m in self.image_contrast_proj.modules():
            if isinstance(m, torch.nn.Linear):
                torch.nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias)
        for param in self.image_contrast_proj.parameters():
            param.requires_grad = True

        self.text_contrast_proj = torch.nn.Sequential(
            torch.nn.Linear(d_llm, contrastive_dim),
            torch.nn.BatchNorm1d(contrastive_dim),
            torch.nn.GELU(),
            torch.nn.Linear(contrastive_dim, contrastive_dim),
            torch.nn.BatchNorm1d(contrastive_dim)
        )   
        for m in self.text_contrast_proj.modules():
            if isinstance(m, torch.nn.Linear):
                torch.nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias)
        for param in self.text_contrast_proj.parameters():
            param.requires_grad = True

        self.contrastive_loss_weight = contrastive_loss_weight
        self.contrastive_genloss_weight = contrastive_genloss_weight

        # 添加对比损失函数
        self.contrastive_loss_fn = MultiModalContrastiveLoss(temperature=0.07) # 默认温度

        # ---------------- 新增：PPO相关组件 ----------------
        # PPO中的价值函数 (Value Head)
        # 将LLM的隐藏状态映射到单一的价值估计
        self.value_head = torch.nn.Sequential(
            torch.nn.Linear(d_llm, d_llm),
            torch.nn.GELU(),
            torch.nn.Linear(d_llm, 1)
        )
        for param in self.value_head.parameters():
            param.requires_grad = True

    def set_temperature(self, temperature):
        """动态设置对比损失函数的温度参数"""
        self.contrastive_loss_fn = MultiModalContrastiveLoss(temperature)

    def set_sim(self,sim):
        self.similarity_threshold = sim

    def forward(self, images, prompts=None, tags=None, is_training=False, target_visual_tokens=128, top_k=None, use_generated_tags=False,
                 ppo_mode=False, input_ids_for_ppo=None, attention_mask_for_ppo=None):
        device = next(self.parameters()).device
        total_loss = None
        gen_loss = None
        contrastive_loss = None
        contrastive_metrics = {}
        
        # ====================== 1. 图像特征提取 ======================
        if isinstance(images, list) and all(isinstance(img, torch.Tensor) for img in images):
            pixel_values = torch.cat(images, dim=0).to(device)
        elif isinstance(images, torch.Tensor):
            pixel_values = images.to(device)
        else:
            clip_inputs = self.processor(images=images, return_tensors="pt").to(device)
            pixel_values = clip_inputs["pixel_values"]

        vision_outputs = self.clip.vision_model(pixel_values=pixel_values)
        vision_hidden_states = vision_outputs.last_hidden_state  # [B, N, D_v]

        image_global_feat = vision_hidden_states.mean(dim=1)  # [B, D_v]
        image_contrast_feat = self.image_contrast_proj(image_global_feat)  # [B, 512]
        image_contrast_feat = F.normalize(image_contrast_feat, p=2, dim=1)

        # ====================== 2. LLM输入处理 ======================
        if target_visual_tokens == 1:
            visual_features_for_llm = vision_hidden_states[:, 0:1, :]
        else:
            visual_features_for_llm = F.adaptive_avg_pool1d(
                vision_hidden_states.permute(0, 2, 1), target_visual_tokens
            ).permute(0, 2, 1)
        multimodal_embed = self.clip2llm_proj(visual_features_for_llm)  # [B, N, D_llm]
        mm_len = multimodal_embed.size(1)

        prompt_text = [
            "严格根据图像内容提取工装设计相关标签(含办公空间、商业空间、酒店空间、商业餐饮等),"
            "需满足:1.优先按照风格空间场景细化,如('中式卧室双人床','北欧风婚纱摄影店','中式特色风格餐馆'),禁止泛化标签(如'客厅','房间');"
            "2. 其次包含'设计类型和功能定位'(如'科技感展厅,互动体验','轻食餐饮,社区社交空间');"
            "3. 接下来是能直接被搜索的工装元素,如('弧形玻璃隔断','透光软膜天花')"
            "4. 最后按'材质,氛围,场景适配'类别顺序排列,同一类别最多生成一个标签(如材质类别可包含'水泥漆,木饰面'),不适用的类别直接省略;"
            "5. 若图片不是场景就仅提取相关物品标签。"
            "6. 仅用英文逗号分隔,不添加任何额外内容（如解释、符号）。仅生成图像中明确存在的标签（如图像有'圆桌包厢'才生成）,不猜测未出现的内容。"
        ] * len(images)
        text_inputs = self.tokenizer(prompt_text, return_tensors="pt", padding=True, truncation=True).to(device)

        llm_embed_layer = self.llm.get_input_embeddings()
        text_embeds = llm_embed_layer(text_inputs.input_ids)  # [B, T, D_llm]
        
        # 如果是PPO模式，我们使用提供的input_ids和attention_mask作为LLM的输入。
        # 否则，我们根据提示和生成/真实标签构建它。
        if ppo_mode and input_ids_for_ppo is not None and attention_mask_for_ppo is not None:
            # 对于PPO，input_ids_for_ppo已经包含了提示词 + 生成的响应
            # 我们需要预置多模态嵌入
            # 假设input_ids_for_ppo对应于生成后的text_inputs
            # 因此我们只需要使用其attention_mask并相应修改嵌入
            # 这部分需要与PPO训练器准备输入的方式仔细对齐。
            # 通常，PPO训练器处理分词并连接提示和响应。
            # 这里我们需要注入我们的多模态嵌入。
            
            # 让我们假设input_ids_for_ppo以提示词tokens开始。
            # 我们需要找到提示词结束的位置来插入多模态嵌入。
            # 对于PPO的前向传播，一个更简单的方法是：让PPO将提示词+响应作为文本输入，
            # 并且我们在它之前注入multimodal_embed。
            
            # 从input_ids_for_ppo重建text_embeds
            # 对于PPO，input_ids_for_ppo将是prompt_text + generated_tags
            # 我们需要将它们分开以插入multimodal_embeds
            
            # 为了简化本例，我们假设`input_ids_for_ppo`包含完整序列（提示+响应）
            # 并且`attention_mask_for_ppo`是其对应的掩码。
            # 我们将生成multimodal_embed并将其预置。
            # 这需要知道input_ids_for_ppo中提示词的长度。
            
            # 在典型的PPO设置中，模型将接收：
            # {
            #    "input_ids": tensor of prompt + response,
            #    "attention_mask": tensor,
            #    "rewards": tensor,
            # }
            # 对我们来说，`input_ids`将是`[多模态token_ids, 提示词token_ids, 响应token_ids]`
            
            # 鉴于当前结构，对于PPO，我们很可能直接将`llm_input_embeds`和`final_attention_mask`
            # 直接输入到LLM中，然后从隐藏状态计算值。
            # `input_ids_for_ppo`和`attention_mask_for_ppo`通常只针对*响应*或*完整序列*。
            # 让我们针对常见的PPO模式进行调整，其中`forward`函数使用`input_ids`和`attention_mask`调用
            # 并且我们需要返回logits和价值估计。

            # 当PPO调用forward时，它将提供`input_ids`和`attention_mask`
            # 它们实际上是`[提示词tokens, 响应tokens]`。
            # 我们需要将视觉tokens*重新注入*到开头。

            # 在这里集成PPO最简单的方法是修改`forward`以满足PPO的特定需求。
            # 当`ppo_mode`为True时：
            #    - `input_ids_for_ppo`将包含分词后的提示词 + 生成的响应。
            #    - `attention_mask_for_ppo`将是其对应的掩码。
            #    - 我们需要通过连接`multimodal_embed`和`text_embeds`（来自`input_ids_for_ppo`）来计算`llm_input_embeds`。
            #    - 我们将返回logits和价值估计。
            
            # 假设input_ids_for_ppo已经连接了提示词 + 生成的文本：
            full_text_embeds = llm_embed_layer(input_ids_for_ppo) # [B, L_text, D_llm]
            llm_input_embeds = torch.cat([multimodal_embed, full_text_embeds], dim=1) # [B, N_mm + L_text, D_llm]
            
            # 调整注意力掩码：将多模态掩码预置到现有文本掩码之前
            mm_attention_mask = torch.ones(multimodal_embed.size(0), mm_len, dtype=torch.long, device=device)
            final_attention_mask = torch.cat([mm_attention_mask, attention_mask_for_ppo], dim=1) # [B, N_mm + L_text]

            llm_outputs = self.llm(
                inputs_embeds=llm_input_embeds,
                attention_mask=final_attention_mask,
                return_dict=True,
                output_hidden_states=True # 需要隐藏状态用于价值头
            )
            
            # 提取最后一个隐藏状态用于价值预测
            last_hidden_state = llm_outputs.hidden_states[-1] # [B, N_mm + L_text, D_llm]
            
            # 仅考虑与生成的文本对应的隐藏状态用于价值头
            # 这意味着切片`last_hidden_state`以排除multimodal_embed和prompt_embeds
            # 假设`input_ids_for_ppo`以提示词开头，然后是响应。
            # 您需要知道`input_ids_for_ppo`中提示词的长度才能正确切片。
            # TRL的PPO中常见的方法是价值头接收整个序列的状态，
            # 但只有*响应*的*最后一个token*用于价值估计。
            
            # 为了简化，我们可以将价值头应用于池化表示或最后一个token的隐藏状态
            # *生成部分*。让我们尝试对*文本部分*（提示词+响应）进行平均
            # 或者简单地取最后一个token的隐藏状态。
            
            # 对于PPO更精确地说：
            # 我们需要用于策略的`logits`和用于价值函数的`value`。
            # `llm_outputs.logits`已经包含了整个序列。
            # 价值头理想情况下应该估计*生成最后一个token*的价值，
            # 或*整个生成序列*的价值。
            # TRL的PPO通常期望每个序列一个标量值。
            
            # 让我们在传递给价值头之前取最后一个token的隐藏状态。
            # 然而，如果是完整的input_ids，最后一个隐藏状态可能来自填充。
            # 最好获取最后一个非填充token的隐藏状态。
            
            sequence_lengths = (final_attention_mask.sum(dim=1) - 1).tolist() # 最后一个token的索引
            # 获取每个序列的最后一个隐藏状态
            last_token_hidden_states = last_hidden_state[torch.arange(last_hidden_state.size(0), device=device), sequence_lengths] # [B, D_llm]
            
            values = self.value_head(last_token_hidden_states).squeeze(-1) # [B]
            
            # 返回PPO训练器的logits和values
            return llm_outputs.logits, values

        else: # 常规训练或推理模式（非PPO）
            llm_input_embeds = torch.cat([multimodal_embed, text_embeds], dim=1)  # [B, N+T, D_llm]
            mm_attention_mask = torch.ones(multimodal_embed.size(0), mm_len, dtype=torch.long, device=device)
            final_attention_mask = torch.cat([mm_attention_mask, text_inputs.attention_mask], dim=1)  # [B, N+T]

            if tags is not None:
                # ---------------- 3.1 生成损失（交叉熵）----------------
                prefix_len = multimodal_embed.size(1) + text_embeds.size(1)
                tag_tokenized = self.tokenizer(
                    [t + self.tokenizer.eos_token for t in tags],
                    return_tensors="pt",
                    padding="longest",
                    truncation=True,
                    max_length=self.llm.config.max_position_embeddings - prefix_len
                ).input_ids.to(device)

                total_len = prefix_len + tag_tokenized.size(1)
                if total_len > self.llm.config.max_position_embeddings:
                    tag_tokenized = tag_tokenized[:, :self.llm.config.max_position_embeddings - prefix_len]

                labels = torch.full((len(images), prefix_len + tag_tokenized.size(1)), -100, dtype=torch.long, device=device)
                for i in range(len(images)):
                    tag_len = (tag_tokenized[i] != self.tokenizer.pad_token_id).sum().item()
                    labels[i, prefix_len : prefix_len + tag_len] = tag_tokenized[i, :tag_len]
                tag_embeds = llm_embed_layer(tag_tokenized)
                llm_input_embeds = torch.cat([llm_input_embeds, tag_embeds], dim=1)
                tag_attention_mask = (tag_tokenized != self.tokenizer.pad_token_id).long().to(device)
                final_attention_mask = torch.cat([final_attention_mask, tag_attention_mask], dim=1)

                llm_outputs = self.llm(
                    inputs_embeds=llm_input_embeds,
                    attention_mask=final_attention_mask,
                    labels=labels,
                    return_dict=True
                )
                gen_loss = llm_outputs.loss.mean()
                total_loss = gen_loss

                # ---------------- 3.2 对比损失（图像-文本特征对比）----------------
                with torch.no_grad():
                    tag_embeds_for_contrast = llm_embed_layer(tag_tokenized) # [B, L, D_llm]
                    tag_mask_for_contrast = (tag_tokenized != self.tokenizer.pad_token_id).unsqueeze(-1).float() # [B, L, 1]
                    tag_mask_sum = tag_mask_for_contrast.sum(dim=1) # [B, 1]
                    tag_mask_sum = torch.clamp(tag_mask_sum, min=1.0)
                    text_global_feat = (tag_embeds_for_contrast * tag_mask_for_contrast).sum(dim=1) / tag_mask_sum # [B, D_llm]
                
                text_contrast_feat = self.text_contrast_proj(text_global_feat) # [B, 512]
                text_contrast_feat = F.normalize(text_contrast_feat, p=2, dim=1)

                contrastive_loss_real, contrastive_metrics = self.contrastive_loss_fn(
                    image_contrast_feat, text_contrast_feat
                )

                contrastive_loss_gen = 0.0
                if use_generated_tags and is_training:
                    generated_tags = self.generate_tags(images, prompts, top_k=5)

                    for i in range(len(generated_tags)):
                        if not generated_tags[i].strip():
                            generated_tags[i] = "默认标签"
                    
                    gen_tag_tokenized = self.tokenizer(
                        [t + self.tokenizer.eos_token for t in generated_tags],
                        return_tensors="pt", padding="longest", truncation=True
                    ).input_ids.to(device)

                    with torch.no_grad():
                        gen_tag_embeds = llm_embed_layer(gen_tag_tokenized)
                        gen_tag_mask = (gen_tag_tokenized != self.tokenizer.pad_token_id).unsqueeze(-1).float()
                        gen_tag_mask_sum = gen_tag_mask.sum(dim=1)
                        gen_tag_mask_sum = torch.clamp(gen_tag_mask_sum, min=1.0)
                        gen_text_global_feat = (gen_tag_embeds * gen_tag_mask).sum(dim=1) / gen_tag_mask_sum
                    
                    gen_text_contrast_feat = self.text_contrast_proj(gen_text_global_feat)
                    gen_text_contrast_feat = F.normalize(gen_text_contrast_feat, p=2, dim=1)

                    contrastive_loss_gen, _ = self.contrastive_loss_fn(
                        image_contrast_feat, gen_text_contrast_feat
                    )

                if use_generated_tags and is_training:
                    contrastive_loss = (1 - self.contrastive_genloss_weight) * contrastive_loss_real + self.contrastive_genloss_weight * contrastive_loss_gen
                else:
                    contrastive_loss = contrastive_loss_real
                
                total_loss = (1 - self.contrastive_loss_weight) * gen_loss + self.contrastive_loss_weight * contrastive_loss

            else: # 标签生成的推理模式
                generated_ids = self.llm.generate(
                    inputs_embeds=llm_input_embeds,
                    attention_mask=final_attention_mask,
                    max_new_tokens=30,
                    num_beams=4,
                    early_stopping=True,
                    pad_token_id=self.tokenizer.pad_token_id,
                    eos_token_id=self.tokenizer.eos_token_id,
                    repetition_penalty=1.5,
                    no_repeat_ngram_size=2
                )

                generated_tags = []
                for ids in generated_ids:
                    ids_list = ids.tolist()
                    if self.tokenizer.eos_token_id in ids_list:
                        eos_idx = ids_list.index(self.tokenizer.eos_token_id)
                        ids_truncated = ids_list[:eos_idx]
                    else:
                        ids_truncated = ids_list
                    
                    tags_text = self.tokenizer.decode(ids_truncated, skip_special_tokens=True)
                    tags_list = [tag.strip() for tag in tags_text.split(',') if tag.strip()]
                    generated_tags.append(tags_list)

                if self.similarity_calculator is not None and not is_training:
                    filtered_tags_all = []
                    filtered_sims_all = []
                    all_sims_all = []

                    for i, tags_list in enumerate(generated_tags):
                        img = images[i] if isinstance(images, list) else images
                        filtered_texts, filtered_sims, all_sims = self.similarity_calculator.filter_texts_by_similarity(
                            img, tags_list, threshold=self.similarity_threshold, verbose=False
                        )
                        if top_k is not None and top_k > 0 and len(filtered_texts) > top_k:
                            filtered_texts = filtered_texts[:top_k]
                            filtered_sims = filtered_sims[:top_k]
                        
                        if not filtered_texts and len(all_sims) > 0:
                            sorted_tags = sorted(zip(tags_list, all_sims), 
                                                 key=lambda x: x[1], reverse=True)[:3]
                            filtered_texts = [tag for tag, _ in sorted_tags]
                            filtered_sims = [sim for _, sim in sorted_tags]
                        
                        filtered_tags_all.append(",".join(filtered_texts))
                        filtered_sims_all.append(filtered_sims)
                        all_sims_all.append(all_sims)

                    return filtered_tags_all, filtered_sims_all, all_sims_all
                else:
                    return [",".join(tags) for tags in generated_tags], [], []
        
        if ppo_mode:
            # 如果在ppo_mode下代码执行到这里，意味着前向传播被调用
            # 用于生成，并且我们已经返回了(logits, values)。
            # 在PPO的主训练循环中不应该到达此'return'，因为它期望
            # 外部的奖励。
            # 代码的这一部分用于“初始”生成，或者如果forward被调用
            # 只是为了根据输入生成响应，然后奖励在外部计算。
            # 我们将需要一个单独的奖励计算机制。
            pass # 这部分由顶部的初始`if ppo_mode`块处理

        # 返回损失和指标（用于非PPO训练）
        return {
            "loss": total_loss,
            "gen_loss": gen_loss,
            "contrastive_loss": contrastive_loss,
            **contrastive_metrics
        }

    def generate_tags(self, images, prompts, top_k=5, max_new_tokens=64, do_sample=False, temperature=1.0):
        """
        生成图像标签。此函数设计为同时支持通用标签生成和 PPO 模式。

        Args:
            images (Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor]): 输入图像。
            prompts (Union[str, List[str]]): 用于生成标签的文本提示。
            top_k (int, optional): 对于采样生成，考虑的最高概率词汇数量。默认为 5。
            max_new_tokens (int, optional): 生成的最大新词汇数量。默认为 64。
            do_sample (bool, optional): 是否使用采样生成。PPO 通常需要采样。默认为 False (使用 beam search)。
            temperature (float, optional): 采样温度。用于控制生成文本的随机性。默认为 1.0。

        Returns:
            Tuple: 包含以下元素的元组：
                - final_decoded_tag_strings (List[str]): 逗号分隔的最终生成标签字符串列表（可能已过滤）。
                - filtered_sims_all (List[List[float]]): 过滤后的标签与图像相似度列表。
                - all_sims_all (List[List[float]]): 所有原始标签与图像相似度列表。
                - full_generated_ids (torch.Tensor): LLM 生成的完整 token ID 序列（包含多模态前缀和提示）。
                - prompt_input_ids_for_ppo (torch.Tensor): 纯文本提示的 token ID，用于 PPO 的 KL 散度计算。
        """
        # 暂时禁用 KV 缓存，因为它可能会干扰梯度检查点（如果使用的话）
        original_use_cache = self.llm.config.use_cache
        self.llm.config.use_cache = False 
        
        device = next(self.parameters()).device

        # 1. 准备 LLM 输入（多模态嵌入 + 提示嵌入）
        # 处理图像输入，将其转换为 pixel_values
        if isinstance(images, list) and all(isinstance(img, torch.Tensor) for img in images):
            pixel_values = torch.cat(images, dim=0).to(device)
        elif isinstance(images, torch.Tensor):
            pixel_values = images.to(device)
        else:
            clip_inputs = self.processor(images=images, return_tensors="pt").to(device)
            pixel_values = clip_inputs["pixel_values"]

        # 通过 CLIP 视觉模型获取视觉特征
        vision_outputs = self.clip.vision_model(pixel_values=pixel_values)
        vision_hidden_states = vision_outputs.last_hidden_state
        
        # 根据 target_visual_tokens 数量，对视觉特征进行自适应池化
        if self.target_visual_tokens == 1:
            visual_features_for_llm = vision_hidden_states[:, 0:1, :]
        else:
            visual_features_for_llm = F.adaptive_avg_pool1d(
                vision_hidden_states.permute(0, 2, 1), self.target_visual_tokens
            ).permute(0, 2, 1)

        # 将 CLIP 视觉特征投影到 LLM 嵌入空间
        multimodal_embed = self.clip2llm_proj(visual_features_for_llm)
        multimodal_prefix_len = multimodal_embed.size(1) # 多模态前缀的长度

        # 准备文本提示输入
        # 确保 prompts 是一个列表，以便批量处理
        if not isinstance(prompts, list):
            prompts = [prompts] * multimodal_embed.size(0) 

        # 对文本提示进行 tokenization，获取其 input_ids 和 attention_mask
        # prompt_input_ids_for_ppo 是纯文本提示的 token ID，用于 PPO 的 KL 散度计算
        text_inputs = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(device)
        prompt_input_ids_for_ppo = text_inputs.input_ids 

        # 获取 LLM 的嵌入层，并将文本 token 转换为嵌入
        llm_embed_layer = self.llm.get_input_embeddings()
        text_embeds = llm_embed_layer(prompt_input_ids_for_ppo)

        # 将多模态嵌入和文本嵌入拼接起来，作为 LLM 的最终输入嵌入
        llm_input_embeds = torch.cat([multimodal_embed, text_embeds], dim=1)
        
        # 构建最终的注意力掩码，结合多模态前缀和文本提示的掩码
        mm_attention_mask = torch.ones(multimodal_embed.size(0), multimodal_prefix_len, dtype=torch.long, device=device)
        final_attention_mask = torch.cat([mm_attention_mask, text_inputs.attention_mask], dim=1)

        # 4. 使用 LLM 生成 token
        with torch.no_grad(): # 在生成过程中通常不计算梯度
            full_generated_ids = self.llm.generate(
                inputs_embeds=llm_input_embeds,
                attention_mask=final_attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=do_sample, # PPO 通常需要采样
                temperature=temperature, # 采样温度
                top_k=top_k if do_sample else None, # 仅在采样时应用 top_k
                num_beams=1 if do_sample else 4, # 采样时 num_beams=1，否则使用 Beam Search
                early_stopping=True,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                repetition_penalty=1.5,
                no_repeat_ngram_size=2
            )
        # 恢复 LLM 的 KV 缓存设置
        self.llm.config.use_cache = original_use_cache

        # 5. 解码 LLM 生成的原始标签（在过滤之前）
        decoded_tags_list_raw = [] # 存储解码后的原始标签列表 (e.g., [['tag1', 'tag2'], ['tagA']])
        
        # 计算 LLM 输入 token 的总长度（多模态前缀长度 + 文本提示 token 长度）
        input_token_len = multimodal_prefix_len + text_inputs.input_ids.size(1)

        for i, ids in enumerate(full_generated_ids):
            # 提取 LLM 实际生成的新 token 部分
            generated_response_ids = ids[input_token_len:]
            
            # 处理生成的 ID，截断到 EOS token 或最大长度
            ids_list = generated_response_ids.tolist()
            if self.tokenizer.eos_token_id in ids_list:
                eos_idx = ids_list.index(self.tokenizer.eos_token_id)
                ids_truncated = ids_list[:eos_idx]
            else:
                ids_truncated = ids_list
            
            # 解码 token 为文本，并分割成标签列表
            tags_text = self.tokenizer.decode(ids_truncated, skip_special_tokens=True)
            tags_list = [tag.strip() for tag in tags_text.split(',') if tag.strip()]
            decoded_tags_list_raw.append(tags_list)

        # 6. 根据 similarity_calculator 是否存在，进行标签过滤处理
        # 初始化最终的标签列表和相似度列表，确保它们在所有情况下都被赋值
        final_decoded_tags_list = [] # 存储最终的标签列表（可能经过过滤）
        filtered_sims_all = [] # 存储过滤后的相似度列表
        all_sims_all = [] # 存储所有原始相似度列表

        if self.similarity_calculator is not None:
            # 如果存在相似度计算器，则进行过滤
            for i, tags_list_for_sample in enumerate(decoded_tags_list_raw):
                # 确保传递给 filter_texts_by_similarity 的是单个图像 tensor
                current_image_tensor = pixel_values[i].unsqueeze(0)

                # 调用相似度计算器进行过滤
                filtered_texts_for_sample, filtered_sims_for_sample, all_sims_for_sample = \
                    self.similarity_calculator.filter_texts_by_similarity(
                        current_image_tensor, tags_list_for_sample, threshold=self.similarity_threshold, verbose=False
                    )
                
                # 应用 top_k 截断（如果过滤后仍然很多）
                if top_k is not None and top_k > 0 and len(filtered_texts_for_sample) > top_k:
                    filtered_texts_for_sample = filtered_texts_for_sample[:top_k]
                    filtered_sims_for_sample = filtered_sims_for_sample[:top_k]
                
                # 如果过滤后没有标签，则从原始标签中选择相似度最高的 3 个作为回退
                if not filtered_texts_for_sample and len(all_sims_for_sample) > 0:
                    sorted_tags_with_sims = sorted(zip(tags_list_for_sample, all_sims_for_sample), 
                                                key=lambda x: x[1], reverse=True)[:3]
                    filtered_texts_for_sample = [tag for tag, _ in sorted_tags_with_sims]
                    filtered_sims_for_sample = [sim for _, sim in sorted_tags_with_sims]
                
                # 将处理后的结果添加到最终列表中
                final_decoded_tags_list.append(filtered_texts_for_sample)
                filtered_sims_all.append(filtered_sims_for_sample)
                all_sims_all.append(all_sims_for_sample)
        else:
            # 如果没有 similarity_calculator，则最终标签列表就是原始解码结果
            final_decoded_tags_list = list(decoded_tags_list_raw)
            # 对应的相似度列表都是空的，但要与批处理大小匹配
            # 为每个样本添加一个空的列表，以保持结构一致性
            filtered_sims_all = [[] for _ in decoded_tags_list_raw]
            all_sims_all = [[] for _ in decoded_tags_list_raw]
        
        # 7. 将最终的标签列表转换为逗号分隔的字符串列表
        final_decoded_tag_strings = [",".join(tags) for tags in final_decoded_tags_list]

        # 返回 PPO 训练器和评估函数所需的参数
        # 统一返回这五个变量，无论是否进行相似度过滤
        return final_decoded_tag_strings, filtered_sims_all, all_sims_all, full_generated_ids, prompt_input_ids_for_ppo
        

# 5. 构建整体模型 (保持不变，但PPO训练时需要额外实例化PPOTrainer)
def build_multimodal_model(config):
    print("[INFO] Loading CLIP + LLM...")
    clip_model, processor = build_clip_model(config.clip_path)
    llm_model, tokenizer = build_llm_model(config.baichuan_path)

    if config.use_lora:
       clip_model, llm_model = add_lora_to_models(clip_model, llm_model)

    if hasattr(config, 'independent_clip_path') and config.independent_clip_path:
        print("[INFO] Loading independent Chinese-CLIP for similarity filtering...")
        similarity_calculator = ChineseCLIPSimilarityCalculator(
            clip_model_path=config.independent_clip_path,
            device="cuda" if torch.cuda.is_available() else "cpu"
        )
    else:
        similarity_calculator = None
        print("[WARNING] No independent CLIP path provided. Similarity filtering will be disabled.")

    model = MultiModalTaggingModel(clip_model, llm_model, processor, tokenizer,
                                   similarity_calculator=similarity_calculator,
                                   contrastive_loss_weight = config.contrastive_loss_weight,
                                   contrastive_genloss_weight = config.contrastive_genloss_weight
                                   )
    return model